"""
高速列车轴承智能故障诊断 - 深度学习模型构建模块

本模块包含：
1. 多种深度学习模型构建（CNN、ResNet、Transformer、混合模型）
2. 模型训练和优化
3. 模型集成和性能评估

作者：数学建模团队
版本：1.0
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
import tensorflow as tf
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.layers import (Dense, Conv1D, Conv2D, MaxPooling1D, MaxPooling2D, 
                                   Flatten, Dropout, BatchNormalization, Activation,
                                   Input, LSTM, Bidirectional, GlobalAveragePooling1D,
                                   Attention, MultiHeadAttention, LayerNormalization)
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau, ModelCheckpoint
from tensorflow.keras.utils import to_categorical
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class DeepLearningModels:
    """深度学习模型构建类"""
    
    def __init__(self, input_shape, num_classes):
        """
        初始化模型构建器
        
        Args:
            input_shape (tuple): 输入形状
            num_classes (int): 类别数量
        """
        self.input_shape = input_shape
        self.num_classes = num_classes
        
    def build_cnn_model(self):
        """构建1D CNN模型"""
        model = Sequential([
            Conv1D(64, 3, activation='relu', input_shape=self.input_shape),
            BatchNormalization(),
            MaxPooling1D(2),
            Conv1D(128, 3, activation='relu'),
            BatchNormalization(),
            MaxPooling1D(2),
            Conv1D(256, 3, activation='relu'),
            BatchNormalization(),
            GlobalAveragePooling1D(),
            Dropout(0.5),
            Dense(128, activation='relu'),
            Dropout(0.3),
            Dense(self.num_classes, activation='softmax')
        ])
        return model
    
    def build_resnet_model(self):
        """构建ResNet模型"""
        def residual_block(x, filters, kernel_size=3):
            shortcut = x
            x = Conv1D(filters, kernel_size, padding='same')(x)
            x = BatchNormalization()(x)
            x = Activation('relu')(x)
            x = Conv1D(filters, kernel_size, padding='same')(x)
            x = BatchNormalization()(x)
            x = tf.keras.layers.add([x, shortcut])
            x = Activation('relu')(x)
            return x
        
        inputs = Input(shape=self.input_shape)
        x = Conv1D(64, 7, padding='same')(inputs)
        x = BatchNormalization()(x)
        x = Activation('relu')(x)
        x = MaxPooling1D(3)(x)
        
        for filters in [64, 128, 256]:
            for _ in range(2):
                x = residual_block(x, filters)
            x = MaxPooling1D(2)(x)
        
        x = GlobalAveragePooling1D()(x)
        x = Dense(256, activation='relu')(x)
        x = Dropout(0.5)(x)
        outputs = Dense(self.num_classes, activation='softmax')(x)
        
        return Model(inputs, outputs)
    
    def build_transformer_model(self):
        """构建Transformer模型"""
        def transformer_encoder(inputs, head_size, num_heads, ff_dim, dropout=0):
            # Attention and Normalization
            x = MultiHeadAttention(key_dim=head_size, num_heads=num_heads, dropout=dropout)(inputs, inputs)
            x = Dropout(dropout)(x)
            x = LayerNormalization(epsilon=1e-6)(x + inputs)
            
            # Feed Forward Part
            y = Dense(ff_dim, activation="relu")(x)
            y = Dense(inputs.shape[-1])(y)
            y = Dropout(dropout)(y)
            y = LayerNormalization(epsilon=1e-6)(x + y)
            return y
        
        inputs = Input(shape=self.input_shape)
        x = inputs
        
        for _ in range(2):
            x = transformer_encoder(x, head_size=256, num_heads=4, ff_dim=512, dropout=0.1)
        
        x = GlobalAveragePooling1D()(x)
        x = Dropout(0.5)(x)
        outputs = Dense(self.num_classes, activation='softmax')(x)
        
        return Model(inputs, outputs)
    
    def build_hybrid_model(self):
        """构建CNN-Transformer混合模型"""
        inputs = Input(shape=self.input_shape)
        
        # CNN部分
        x = Conv1D(64, 3, activation='relu')(inputs)
        x = BatchNormalization()(x)
        x = MaxPooling1D(2)(x)
        x = Conv1D(128, 3, activation='relu')(x)
        x = BatchNormalization()(x)
        x = MaxPooling1D(2)(x)
        
        # Transformer部分
        x = MultiHeadAttention(num_heads=4, key_dim=64)(x, x)
        x = LayerNormalization(epsilon=1e-6)(x)
        x = Dropout(0.1)(x)
        
        x = GlobalAveragePooling1D()(x)
        x = Dense(256, activation='relu')(x)
        x = Dropout(0.5)(x)
        outputs = Dense(self.num_classes, activation='softmax')(x)
        
        return Model(inputs, outputs)
    
    def build_lstm_model(self):
        """构建LSTM模型"""
        model = Sequential([
            LSTM(128, return_sequences=True, input_shape=self.input_shape),
            Dropout(0.2),
            LSTM(64, return_sequences=True),
            Dropout(0.2),
            LSTM(32),
            Dropout(0.2),
            Dense(64, activation='relu'),
            Dropout(0.3),
            Dense(self.num_classes, activation='softmax')
        ])
        return model
    
    def build_bidirectional_lstm_model(self):
        """构建双向LSTM模型"""
        model = Sequential([
            Bidirectional(LSTM(128, return_sequences=True), input_shape=self.input_shape),
            Dropout(0.2),
            Bidirectional(LSTM(64, return_sequences=True)),
            Dropout(0.2),
            Bidirectional(LSTM(32)),
            Dropout(0.2),
            Dense(64, activation='relu'),
            Dropout(0.3),
            Dense(self.num_classes, activation='softmax')
        ])
        return model


class ModelTrainer:
    """模型训练类"""
    
    def __init__(self, model, model_name):
        """
        初始化模型训练器
        
        Args:
            model: 要训练的模型
            model_name (str): 模型名称
        """
        self.model = model
        self.model_name = model_name
        self.history = None
        
    def compile_model(self, learning_rate=0.001):
        """
        编译模型
        
        Args:
            learning_rate (float): 学习率
        """
        self.model.compile(
            optimizer=Adam(learning_rate=learning_rate),
            loss='categorical_crossentropy',
            metrics=['accuracy', 'precision', 'recall']
        )
    
    def train_model(self, X_train, y_train, X_val, y_val, epochs=100, batch_size=32):
        """
        训练模型
        
        Args:
            X_train (np.array): 训练特征
            y_train (np.array): 训练标签
            X_val (np.array): 验证特征
            y_val (np.array): 验证标签
            epochs (int): 训练轮数
            batch_size (int): 批次大小
        """
        callbacks = [
            EarlyStopping(monitor='val_loss', patience=15, restore_best_weights=True),
            ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5, min_lr=1e-7),
            ModelCheckpoint(f'best_{self.model_name}.h5', save_best_only=True)
        ]
        
        self.history = self.model.fit(
            X_train, y_train,
            validation_data=(X_val, y_val),
            epochs=epochs,
            batch_size=batch_size,
            callbacks=callbacks,
            verbose=1
        )
    
    def evaluate_model(self, X_test, y_test):
        """
        评估模型
        
        Args:
            X_test (np.array): 测试特征
            y_test (np.array): 测试标签
            
        Returns:
            dict: 评估指标
        """
        results = self.model.evaluate(X_test, y_test, verbose=0)
        metrics = dict(zip(self.model.metrics_names, results))
        return metrics
    
    def plot_training_history(self):
        """绘制训练历史"""
        if self.history is None:
            print("模型尚未训练，无法绘制训练历史")
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        
        ax1.plot(self.history.history['accuracy'], label='训练准确率')
        ax1.plot(self.history.history['val_accuracy'], label='验证准确率')
        ax1.set_title('模型准确率')
        ax1.set_xlabel('轮次')
        ax1.set_ylabel('准确率')
        ax1.legend()
        
        ax2.plot(self.history.history['loss'], label='训练损失')
        ax2.plot(self.history.history['val_loss'], label='验证损失')
        ax2.set_title('模型损失')
        ax2.set_xlabel('轮次')
        ax2.set_ylabel('损失')
        ax2.legend()
        
        plt.tight_layout()
        plt.show()
    
    def predict(self, X):
        """
        预测
        
        Args:
            X (np.array): 输入数据
            
        Returns:
            np.array: 预测结果
        """
        return self.model.predict(X)
    
    def predict_classes(self, X):
        """
        预测类别
        
        Args:
            X (np.array): 输入数据
            
        Returns:
            np.array: 预测类别
        """
        predictions = self.predict(X)
        return np.argmax(predictions, axis=1)


class ModelEnsemble:
    """模型集成类"""
    
    def __init__(self, models):
        """
        初始化模型集成器
        
        Args:
            models (list): 模型列表
        """
        self.models = models
        
    def predict(self, X):
        """
        集成预测
        
        Args:
            X (np.array): 输入数据
            
        Returns:
            np.array: 平均预测概率
        """
        predictions = []
        for model in self.models:
            pred = model.predict(X)
            predictions.append(pred)
        
        # 平均预测概率
        avg_pred = np.mean(predictions, axis=0)
        return avg_pred
    
    def predict_classes(self, X):
        """
        预测类别
        
        Args:
            X (np.array): 输入数据
            
        Returns:
            np.array: 预测类别
        """
        predictions = self.predict(X)
        return np.argmax(predictions, axis=1)
    
    def evaluate_ensemble(self, X_test, y_test, class_names=None):
        """
        评估集成模型
        
        Args:
            X_test (np.array): 测试特征
            y_test (np.array): 测试标签
            class_names (list): 类别名称
        """
        y_pred_proba = self.predict(X_test)
        y_pred = np.argmax(y_pred_proba, axis=1)
        y_true = np.argmax(y_test, axis=1) if y_test.ndim > 1 else y_test
        
        print("分类报告:")
        print(classification_report(y_true, y_pred, target_names=class_names))
        
        # 绘制混淆矩阵
        plt.figure(figsize=(10, 8))
        cm = confusion_matrix(y_true, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                   xticklabels=class_names, yticklabels=class_names)
        plt.title('混淆矩阵')
        plt.xlabel('预测标签')
        plt.ylabel('真实标签')
        plt.show()
        
        return {
            'accuracy': accuracy_score(y_true, y_pred),
            'predictions': y_pred,
            'probabilities': y_pred_proba
        }


class ModelComparison:
    """模型比较类"""
    
    def __init__(self, models_dict, X_test, y_test):
        """
        初始化模型比较器
        
        Args:
            models_dict (dict): 模型字典 {name: model}
            X_test (np.array): 测试特征
            y_test (np.array): 测试标签
        """
        self.models = models_dict
        self.X_test = X_test
        self.y_test = y_test
        self.results = {}
        
    def compare_models(self):
        """比较所有模型"""
        for name, model in self.models.items():
            print(f"\n评估模型: {name}")
            
            # 预测
            y_pred = model.predict_classes(self.X_test)
            y_true = np.argmax(self.y_test, axis=1) if self.y_test.ndim > 1 else self.y_test
            
            # 计算指标
            accuracy = accuracy_score(y_true, y_pred)
            
            self.results[name] = {
                'accuracy': accuracy,
                'predictions': y_pred
            }
            
            print(f"准确率: {accuracy:.4f}")
        
        return self.results
    
    def plot_comparison(self):
        """绘制模型比较图"""
        if not self.results:
            print("请先运行 compare_models()")
            return
        
        # 提取准确率
        model_names = list(self.results.keys())
        accuracies = [self.results[name]['accuracy'] for name in model_names]
        
        # 绘制柱状图
        plt.figure(figsize=(12, 6))
        bars = plt.bar(model_names, accuracies, color='skyblue', alpha=0.7)
        plt.title('模型性能比较')
        plt.xlabel('模型')
        plt.ylabel('准确率')
        plt.ylim(0, 1)
        
        # 在柱子上添加数值
        for bar, acc in zip(bars, accuracies):
            plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                    f'{acc:.3f}', ha='center', va='bottom')
        
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.show()


class ModelOptimizer:
    """模型优化类"""
    
    def __init__(self, model_builder, input_shape, num_classes):
        """
        初始化模型优化器
        
        Args:
            model_builder: 模型构建函数
            input_shape (tuple): 输入形状
            num_classes (int): 类别数量
        """
        self.model_builder = model_builder
        self.input_shape = input_shape
        self.num_classes = num_classes
        
    def hyperparameter_search(self, X_train, y_train, X_val, y_val, 
                            param_grid, max_trials=10):
        """
        超参数搜索
        
        Args:
            X_train (np.array): 训练特征
            y_train (np.array): 训练标签
            X_val (np.array): 验证特征
            y_val (np.array): 验证标签
            param_grid (dict): 参数网格
            max_trials (int): 最大试验次数
            
        Returns:
            dict: 最佳参数和结果
        """
        best_score = 0
        best_params = None
        results = []
        
        for trial in range(max_trials):
            # 随机选择参数
            params = {}
            for param, values in param_grid.items():
                params[param] = np.random.choice(values)
            
            # 构建和训练模型
            model = self.model_builder(self.input_shape, self.num_classes)
            trainer = ModelTrainer(model, f'trial_{trial}')
            trainer.compile_model(learning_rate=params.get('learning_rate', 0.001))
            
            # 训练模型
            trainer.train_model(X_train, y_train, X_val, y_val, 
                              epochs=params.get('epochs', 50),
                              batch_size=params.get('batch_size', 32))
            
            # 评估模型
            metrics = trainer.evaluate_model(X_val, y_val)
            score = metrics['accuracy']
            
            results.append({
                'trial': trial,
                'params': params,
                'score': score
            })
            
            if score > best_score:
                best_score = score
                best_params = params
            
            print(f"试验 {trial+1}: 准确率 = {score:.4f}")
        
        return {
            'best_params': best_params,
            'best_score': best_score,
            'all_results': results
        }


def main():
    """主函数示例"""
    # 示例数据
    np.random.seed(42)
    X = np.random.randn(1000, 30, 1)
    y = np.random.randint(0, 4, 1000)
    y_cat = to_categorical(y)
    
    # 划分数据集
    X_train, X_test, y_train, y_test = train_test_split(
        X, y_cat, test_size=0.2, random_state=42, stratify=y
    )
    X_train, X_val, y_train, y_val = train_test_split(
        X_train, y_train, test_size=0.2, random_state=42
    )
    
    # 模型参数
    input_shape = (30, 1)
    num_classes = 4
    class_names = ['正常', '外圈故障', '内圈故障', '滚动体故障']
    
    # 构建模型
    model_builder = DeepLearningModels(input_shape, num_classes)
    
    # 训练多个模型
    models = {}
    
    # CNN模型
    cnn_model = model_builder.build_cnn_model()
    cnn_trainer = ModelTrainer(cnn_model, 'CNN')
    cnn_trainer.compile_model()
    cnn_trainer.train_model(X_train, y_train, X_val, y_val, epochs=20)
    models['CNN'] = cnn_trainer
    
    # ResNet模型
    resnet_model = model_builder.build_resnet_model()
    resnet_trainer = ModelTrainer(resnet_model, 'ResNet')
    resnet_trainer.compile_model()
    resnet_trainer.train_model(X_train, y_train, X_val, y_val, epochs=20)
    models['ResNet'] = resnet_trainer
    
    # Transformer模型
    transformer_model = model_builder.build_transformer_model()
    transformer_trainer = ModelTrainer(transformer_model, 'Transformer')
    transformer_trainer.compile_model()
    transformer_trainer.train_model(X_train, y_train, X_val, y_val, epochs=20)
    models['Transformer'] = transformer_trainer
    
    # 模型比较
    comparison = ModelComparison(models, X_test, y_test)
    results = comparison.compare_models()
    comparison.plot_comparison()
    
    # 集成模型
    ensemble_models = [cnn_model, resnet_model, transformer_model]
    ensemble = ModelEnsemble(ensemble_models)
    ensemble_results = ensemble.evaluate_ensemble(X_test, y_test, class_names)
    
    print(f"\n集成模型准确率: {ensemble_results['accuracy']:.4f}")


if __name__ == "__main__":
    main()
